import pandas as pd
import os
import sys
import re


"""

执行时间：
执行 
c-lop-analyze.py 
-> get_all_real_total_time_and_predict_time_to_summary_csv.py 
-> get_avg_precision_from_summary_csv.py 
之后进行排序输出

该脚本作用：
传入一个csv文件，根据csv文件中的node 和 proc 进行排序

"""

def extract_number(text):
    """
    Extract the first number from a string.
    
    Args:
        text (str): Input string containing numbers
    
    Returns:
        int: The first number found in the string, or 0 if no number found
    """
    if pd.isna(text):
        return 0
    
    # Use regex to find the first number in the string
    match = re.search(r'\d+', str(text))
    if match:
        return int(match.group())
    return 0

def calculate_relative_error(df):
    """
    计算相对误差：|avg_total_real_comm_time - avg_predict_total_comm_time| / avg_total_real_comm_time * 100
    
    Args:
        df (pd.DataFrame): 包含数据的DataFrame
    
    Returns:
        pd.DataFrame: 添加了relative_error列的DataFrame
    """
    if 'avg_total_real_comm_time' in df.columns and 'avg_predict_total_comm_time' in df.columns:
        # 计算相对误差
        df['relative_error'] = (abs(df['avg_total_real_comm_time'] - df['avg_predict_total_comm_time']) / df['avg_total_real_comm_time']) * 100
        print("已添加relative_error列")
    else:
        print("警告: 缺少计算相对误差所需的列")
    
    return df

def sort_csv_by_node_and_proc(csv_path, output_path=None):
    """
    根据'node_dir'和'proc_dir'列中的数值对CSV文件进行排序。
    先按节点数排序，再按进程数排序。
    同时计算相对误差。
    
    Args:
        csv_path (str): 输入CSV文件的路径
        output_path (str, optional): 保存排序后CSV文件的路径。如果为None，将覆盖输入文件。
    
    Returns:
        pd.DataFrame: 排序后的DataFrame
    """
    try:
        # Read the CSV file
        df = pd.read_csv(csv_path)
        
        # Check if 'node_dir' and 'proc_dir' columns exist
        if 'node_dir' not in df.columns or 'proc_dir' not in df.columns:
            print(f"Error: CSV file must contain 'node_dir' and 'proc_dir' columns.")
            print(f"Available columns: {df.columns.tolist()}")
            return None
        
        print(f"原始数据行数: {len(df)}")
        print(f"排序前前5行:")
        print(df[['node_dir', 'proc_dir']].head())
        
        # Create temporary columns with extracted numeric values for sorting
        df['node_num'] = df['node_dir'].apply(extract_number)
        df['proc_num'] = df['proc_dir'].apply(extract_number)
        
        print(f"提取的数值:")
        print(df[['node_dir', 'node_num', 'proc_dir', 'proc_num']].head())
        
        # Sort the DataFrame by the numeric values (先按节点数，再按进程数)
        sorted_df = df.sort_values(by=['node_num', 'proc_num'])
        
        print(f"排序后前5行:")
        print(sorted_df[['node_dir', 'node_num', 'proc_dir', 'proc_num']].head())
        
        # Remove the temporary columns
        # Remove the temporary columns
        sorted_df = sorted_df.drop(['node_num', 'proc_num'], axis=1)
        
        # 计算相对误差
        sorted_df = calculate_relative_error(sorted_df)
        
        # 确保node_dir列在proc_dir列前面
        columns = sorted_df.columns.tolist()
        if 'node_dir' in columns and 'proc_dir' in columns:
            # 找到node_dir和proc_dir的位置
            node_idx = columns.index('node_dir')
            proc_idx = columns.index('proc_dir')
            
            # 如果proc_dir在node_dir前面，需要调整顺序
            if proc_idx < node_idx:
                # 创建新的列顺序，将node_dir放在proc_dir前面
                new_columns = columns.copy()
                new_columns[proc_idx] = 'node_dir'
                new_columns[node_idx] = 'proc_dir'
                
                # 重新排列DataFrame的列
                sorted_df = sorted_df[new_columns]
                
                print(f"列顺序已调整为: {new_columns}")
        
        # Save the sorted DataFrame
        if output_path:
            sorted_df.to_csv(output_path, index=False)
            print(f"Sorted CSV saved to: {output_path}")
        else:
            sorted_df.to_csv(csv_path, index=False)
            print(f"Original CSV file has been overwritten with sorted data.")
        
        return sorted_df
    
    except Exception as e:
        print(f"Error processing CSV file: {e}")
        return None

if __name__ == "__main__":
   
 
    base_dir_csv = r"F:\PostGraduate\Point-to-Point-Code\C-lop-Prediction\analysis_strong_scaling\analysis_for_all_comm_size_from1-16node_vars\16-128node\all_predict_precision\avg_precision_summary.csv"



    # Check if the file exists
    if not os.path.exists(base_dir_csv):
        print(f"Error: File not found: {base_dir_csv}")
        sys.exit(1)
    
    # Sort the CSV file
    sort_csv_by_node_and_proc(base_dir_csv)
    print("CSV sorting completed successfully.")